package com.numericalmethod.algoquant.model.ralph2009;

import java.util.Arrays;
import java.util.Collection;
import java.util.HashMap;
import java.util.LinkedList;

import org.apache.commons.math3.stat.StatUtils;

import com.numericalmethod.algoquant.model.ralph2009.market.MarketState;
import com.numericalmethod.algoquant.model.ralph2009.market.Ralph2009MarketSDE;
import com.numericalmethod.suanshu.matrix.doubles.Matrix;
import com.numericalmethod.suanshu.stats.stochasticprocess.multivariate.random.MultivariateRandomRealizationOfRandomProcess;
import com.numericalmethod.suanshu.stats.stochasticprocess.timegrid.EvenlySpacedGrid;
import com.numericalmethod.suanshu.stats.timeseries.datastructure.multivariate.realtime.MultivariateRealization;
import com.numericalmethod.suanshu.vector.doubles.dense.DenseVector;

/**
 * 
 * Implementation of the portfolio weight optimization algorithm published by the paper
 * 
 * Momentum and Mean Reversion in Strategic Asset Allocation
 * 
 * (Ralph et al, 2009)
 * 
 * and
 * 
 * A Simulation Approach to Dynamic Portfolio Choice with an
 * Application to Learning About Return Predictability
 * 
 * (Brandt et al. 2005)
 * 
 * @author Paul/Clement/Stephen
 *
 */
public class Ralph2009PortfolioOptimizer {

	private Ralph2009MarketSDE sde;
	private Ralph2009SimulationParameters simParam;
	private double investmentHorizon;
	// number of steps we need to run the simulation (our time step index starts from 0)
	private int numberOfTimeSteps;
	// the shortsell/leverage constraint details
	private Ralph2009PortfolioConstraint constraint;
	
	// event listeners to receive updates
	private Collection<Ralph2009OptimizerStatusListener> listeners;
	
	public Ralph2009PortfolioOptimizer() {
	}
	

	/**
	 * Create an instance of Ralph2009PortfolioOptimizer.
	 * 
	 * @param sde The model SDE of the Ralph 2009 paper.
	 * @param simParam The parameter used in the simulation.
	 * @param investmentHorizon The total investment horizon, measure in the unit of years.
	 * @param constraint The portfolio optimization constraint.
	 */
	public Ralph2009PortfolioOptimizer(Ralph2009MarketSDE sde, 
			Ralph2009SimulationParameters simParam, double investmentHorizon, 
			Ralph2009PortfolioConstraint constraint) {
		this.sde=sde;
		this.simParam=simParam;
		this.investmentHorizon=investmentHorizon;
		numberOfTimeSteps=(int) Math.ceil(investmentHorizon/simParam.getDt());
		this.constraint=constraint;
	}


	/**
	 * 
	 * Return the optimal weight that should allocated to the risky asset
	 * based on the input current market state (St, Mt, Xt), current risk free rate
	 * and the CRRA gamma.
	 * 
	 * @param currentMarketState The current market condition including St: current stock price, Mt: current momentum, Xt: current excess log dividend yield
	 * @param riskFreeRate The risk free rate input to the model.
	 * @param gamma The CRRA Utility function parameter.
	 * 
	 * @return The optimal portfolio weight depends on the user preferring dynamic (i.e. make uses of div information), or myopic (i.e. standard mean/variance portfolio)
	 */
	public Ralph2009OptimizedPortfolio findOptimalWeights(MarketState currentMarketState, double riskFreeRate, double gamma) {
		
		
		// TODO:
		//
		// (1) Decouple Brandt2005 algorithm (multi-dimension version) as an API from here and implement Ralph2009 optimization by invoking the API.
		// (2) Decouple the case of Dynamic Weights and Myopic weights allocation.
		// (3) Generalize the case for Portfolio weighting constraints.
		
		
		// we simulate all the paths,
		LinkedList<MarketState[]> allPaths=simulateAllPaths(currentMarketState);


		// Then we move backward in time and compute the cross paths regression
		// and thus solve the optimization problem described in 
		// 
		// Ralph et al. 2009, 
		// Momentum and Mean Reversion
		// in Strategic Asset Allocation
		//
		// and Brandt et al. 2005 
		// A Simulation Approach to Dynamic Portfolio Choice with an
		// Application to Learning About Return Predictability

		int pathsCount=simParam.getPathsCount();
		
		// all these intermediate value holders have "dynamic weight version"
		// and "myopic weight version"
		// 
		// described in equation (13) of Brandt et al. 2005
		
		// A and B are the two vector described in 
		// equation (18) of Brandt 2005.
		// we have dynamic allocation case, and myopic allocation case.
		double[] simulatedATPlus1Dynamic = new double[pathsCount];
		double[] simulatedBTPlus1Dynamic = new double[pathsCount];
		double[] simulatedATPlus1Myopic = new double[pathsCount];
		double[] simulatedBTPlus1Myopic = new double[pathsCount];
		
		// for holding the regression result of E(a(t+1)) and E(b(t+1)) described in equation (18) of Brandt et al. 2005
		double[] estimatedATPlus1Dynamic = null;
		double[] estimatedBTPlus1Dynamic = null;
		double[] estimatedATPlus1Myopic = null;
		double[] estimatedBTPlus1Myopic = null;			
		
		//initialization for t=T-1
		// these vector represents the psi(t+1) factor 
		// in eq (5) in Brandt et al. 2005
		double[] expectedPsiTPlus1 = new double [pathsCount];
		Arrays.fill(expectedPsiTPlus1, 1.0); //set to 1 for t=T-1
		
		// Rf in equation(5) in Brandt et al. 2005
		double riskFreeReturn = riskFreeRate*simParam.getDt();

    	// construct the result holder.
		Ralph2009OptimizedPortfolio optimizedPortfolio=new Ralph2009OptimizedPortfolio();
		optimizedPortfolio.setConstraint(constraint);
		
		// T-1 is the end of investment horizon 

		// remove the T-1 state , we process from T-2
		MarketState[] stateTPlus1 = allPaths.removeLast();
		double[] lastExcessReturnTPlus1 = null;
		
		// the small t goes from T-2 to 0,
		for (int t=numberOfTimeSteps-2;t>=0;t--) {
			MarketState[] stateT = allPaths.removeLast();
			
			// Re(t+1) of equation (5) in Brandt et al. 2005 
			double[] excessReturnTPlus1 = getExcessReturns(stateT, stateTPlus1, riskFreeReturn);
			
			
			// Z(t+1) of equation (5) in Brandt et al. 2005, 
			double[] zTPlus1 = getExcessDividendYields(stateTPlus1);

			//determine function psi(t+1) and stored in factorVec using eq (5) in Brandt et al. 2005
			
			// NOTE: t is moving backward 
			// and we only need to compute this if we are at the 2nd last time step or earlier time step.
			if (t < numberOfTimeSteps - 2) {
				expectedPsiTPlus1 = getPSITPlus1Observations(expectedPsiTPlus1, optimizedPortfolio.getDynamicWeightAt(t+1), 
						lastExcessReturnTPlus1, riskFreeReturn, gamma);
			}

			// use regression as expectation on Psi(t+1) at time t.
			expectedPsiTPlus1 =  Brandt2005Regression.estimateYValues(excessReturnTPlus1, zTPlus1, expectedPsiTPlus1);
			
			// gather all a(t+1) and b(t+1) from simulation
			// described in equation (12) and (13) in Brandt et al. 2005
			for(int i=0;i<pathsCount;i++){
				// eq (12) of Brandt et al. 2005
				// redefine At+1 = E[psi(Zt+1)*Ret+1] and Bt+1 =E[psi(Zt+1)(Ret+1*Ret+1)]
				// This is the case when innovation of state variable Zt+1 coupled to the optimal weight determination
				simulatedATPlus1Dynamic[i] = expectedPsiTPlus1[i]*excessReturnTPlus1[i]; 
				simulatedBTPlus1Dynamic[i] = expectedPsiTPlus1[i]*(excessReturnTPlus1[i]*excessReturnTPlus1[i]);		
				
				// eq (13) of Brandt et al. 2005
				// if Zt+1 is independent of Ret+1 => reduce to myopic strategy 
				// i.e. purely mean/variance approach, we ignore all external factor Z(t+1)
				simulatedATPlus1Myopic[i] = excessReturnTPlus1[i];  
				simulatedBTPlus1Myopic[i] = (excessReturnTPlus1[i]*excessReturnTPlus1[i]);
			}

			// eq (19) and (20) of Brandt et al. 2005
			// Estimate using cross path regression on the simulated a(t+1) and b(t+1) 

			// dynamic weighting i.e. taking the external factor Z(t+1) into consideration
			estimatedATPlus1Dynamic = Brandt2005Regression.estimateYValues(excessReturnTPlus1, zTPlus1, simulatedATPlus1Dynamic);
			estimatedBTPlus1Dynamic = Brandt2005Regression.estimateYValues(excessReturnTPlus1, zTPlus1, simulatedBTPlus1Dynamic);

			// myopic weighting i.e. assume external factor Z(t+1) is independent to the stock price process.
			estimatedATPlus1Myopic = Brandt2005Regression.estimateYValues(excessReturnTPlus1, zTPlus1, simulatedATPlus1Myopic);
			estimatedBTPlus1Myopic = Brandt2005Regression.estimateYValues(excessReturnTPlus1, zTPlus1, simulatedBTPlus1Myopic);

			
			// finally we compute the optimal weight for this time step t.
			// dynamic weights
			// eq (12) of Brandt et al. 2005
			double optimalDynamicWeight = 
					(1+riskFreeReturn)*StatUtils.mean(estimatedATPlus1Dynamic) / StatUtils.mean(estimatedBTPlus1Dynamic)/gamma;

			
			// myopic weights
			// eq (13) of Brandt et al. 2005
			double optimalMyopicWeight = 
					(1+riskFreeReturn)*StatUtils.mean(estimatedATPlus1Myopic) / StatUtils.mean(estimatedBTPlus1Myopic)/gamma;
			
			// save the result to the portfolio object
			// (constriant on leverage and shortselling are handled inside the portfolio object)
			optimizedPortfolio.addDynamicWeightResult(t, optimalDynamicWeight);
			optimizedPortfolio.addMyopicWeightResult(t, optimalMyopicWeight);
			double corr = getCorrelation(expectedPsiTPlus1,excessReturnTPlus1,simulatedATPlus1Dynamic);
			optimizedPortfolio.addCorrelation(t, corr);
			
			// notify the listeners
			fireWeightComputedAtTimeStep(t, optimizedPortfolio);

			// copy the current state to the state t+1, for next loop processing.
			stateTPlus1=stateT;
			
			// stored for next t step psiT+1 calcuation
			lastExcessReturnTPlus1 = excessReturnTPlus1;
		}
		
		return optimizedPortfolio;
		
	}
	
	
	/**
	 * Register a particular listener to this optimizer to receive an update of status.
	 */
	public void addStatusListener(Ralph2009OptimizerStatusListener l) {
		// create this listeners on demand
		if (listeners==null)
			listeners=new LinkedList<Ralph2009OptimizerStatusListener>();
		listeners.add(l);
	}
	
	/**
	 * Remove a particular listener to this optimizer to NOT receive an update of status.
	 */
	public void removeStatusListener(Ralph2009OptimizerStatusListener l) {
		listeners.remove(l);
		if (listeners.isEmpty())
			listeners=null;
	}

	/**
	 * Helper function to notify the listeners to receive status updates.
	 */
	private void fireWeightComputedAtTimeStep(int timeStepFinished, Ralph2009OptimizedPortfolio p) {
		// do an extra check on listeners collection for performance reason
		// in case no one listening we don't want to spawn extra iterables.
		if (listeners!=null) {
			for (Ralph2009OptimizerStatusListener l : listeners) 
				l.weightAtTimeIndexComputed(timeStepFinished, numberOfTimeSteps, p);
		}
	}
	
	private void fireSimulationCompleted(int pathsSimulated) {
		if (listeners!=null) {
			for (Ralph2009OptimizerStatusListener l : listeners) 
				l.pathsSimulated(pathsSimulated, simParam.getPathsCount());
		}
	}
	
	
	/**
	 * Helper function to simulate all the paths by the given market state 
	 * (price, momentum, excess div yield)  i.e. the (St, Mt, Xt)
	 * 
	 * @return List of market states. Each node represents one time step from zero to terminalTimeIndex. 
	 * The MarketState array represents all paths states at the particular time step. 
	 */
	private LinkedList<MarketState[]> simulateAllPaths(MarketState currentMarketState) {
		
		EvenlySpacedGrid timeGrid=new EvenlySpacedGrid(0.0, investmentHorizon, numberOfTimeSteps+1);
		MultivariateRandomRealizationOfRandomProcess gen = 
				new MultivariateRandomRealizationOfRandomProcess(sde, timeGrid, 
						new DenseVector(currentMarketState.getAssetPrice(), 
								currentMarketState.getMomentum(), 
								currentMarketState.getExcessDividendYield()));
		gen.seed(simParam.getSeed1(), simParam.getSeed2());
		
		int pathsCount=simParam.getPathsCount();

		HashMap<Integer, Matrix> allRealizations=new HashMap<Integer, Matrix>();
		for (int i=0;i<pathsCount;i++) {
			MultivariateRealization realization = gen.nextRealization();
			// realization expected to be (numberOfTimeSteps+1) x 3 matrix
			allRealizations.put(i, realization.toMatrix());
			
			// notify listener
			fireSimulationCompleted(i+1);
		}


		// NOTE: we need to reorganize the data structure instead of keeping it
		// in the original matrix form in order to support efficient
		// cross paths regression during the simulation procedure.
		//
		// TODO: 
		// This is a slow operation and twice the memory consumption if the number of paths is very large.
		// Potentially we can remove the dependency on the MarketState object
		// and directly look up the realization matrix instead. 
		// That's a drastic code changes in the whole algorithm code base 
		// which we cannot do it so quickly right before project submission.
		
		// each element in the list represents all market state paths simulated at time t.
		LinkedList<MarketState[]> allPaths=new LinkedList<MarketState[]>();
		for (int t=0;t<numberOfTimeSteps+1;t++) {
			MarketState[] stateAtT=new MarketState [pathsCount];
			for (int i=0;i<pathsCount;i++) {
				// get the i-th realization
				Matrix m=allRealizations.get(i);
				// in Ralph2009MarketSDE, the first element is St, the second is Mt, and the 3rd is Xt
				double assetPrice = m.get(t+1, 1);
				double momentum = m.get(t+1, 2);
				double excessDivYield = m.get(t+1, 3);
				
				stateAtT[i]=new MarketState(assetPrice, momentum, excessDivYield);
			}
			allPaths.add(stateAtT);
		}
		
		
		return allPaths;
	}
	
	/**
	 * Helper function to compute across all the market state and compute an array of excess return value.
	 */
	private static double[] getExcessReturns(MarketState[] currentState, MarketState[] nextState, double riskFreeReturn) {
		
		double[] allExcessReturns=new double [currentState.length];
		for (int i=0;i<allExcessReturns.length;i++) {
			
			double sTPlus1 = nextState[i].getAssetPrice();
			double sT = currentState[i].getAssetPrice();
			double assetReturn = (sTPlus1 - sT) / sT;
			
			// represents Re(s+1) on equation (5) of Brandt et al. 2005
			double excessReturn = assetReturn - riskFreeReturn; 
			
			allExcessReturns[i]=excessReturn;
		}
		return allExcessReturns;
	}
	
	/**
	 * Helper function to aggregate all Excess Log Dividend Yield in the current market state
	 */
	private static double[] getExcessDividendYields(MarketState[] currentState) {
		double[] result=new double [currentState.length];
		for (int i=0;i<currentState.length;i++) {
			result[i]=currentState[i].getExcessDividendYield();
		}
		return result;
	}
	
	/**
	 * Helper function to compute the psi(Zt+1) for each simulated path in eq (5) of (Brandt 2005).
	 * 
	 * TODO: Refactor this, it is hard to follow the relation of lastFactorArray and the result array 
	 * 
	 * 
	 * @param lastPsiTPlus1Array The psi(Zt+1) computed up to last loop
	 * @param optimalWeightT The x(s) in equation 5. 
	 * @param excessReturnTPlus1 The vector Re(s+1) in equation 5.
	 * @param riskFreeReturn The Rf in equation 5.
	 * @param gamma The CRRA Utility function parameter.
	 * 
	 * @return The psi(Zt+1) multiplication factor for each of the simulation paths at a particular time step. 
	 */
	private static double[] getPSITPlus1Observations(double[] lastPsiTPlus1Array, double optimalWeightT, double[] excessReturnTPlus1, double riskFreeReturn, double gamma) {
		
		double[] newPsiTPlus1Array = new double [excessReturnTPlus1.length];
		
		//determine function psi(t+1)  eqt.5 p.7 in brandt(2005)
		for(int i=0;i<excessReturnTPlus1.length;i++){
			double factor = (optimalWeightT * excessReturnTPlus1[i]) + (1+riskFreeReturn);
			factor = Math.pow(factor,1-gamma);
			newPsiTPlus1Array[i] = lastPsiTPlus1Array[i] * factor;
		}
		
		return newPsiTPlus1Array;
		
	}

	/**
	 * Helper function to compute the correlation coefficient between 2 time series.
	 * 
	 * 
	 * @param xSeries 
	 * @param ySeries
	 * @param xySeries the series of the x*y product  
	 * 
	 * @return The correlation coefficent which is between -1.0 and 1.0
	 */
	private static double getCorrelation(double[] xSeries, double[] ySeries, double[] xySeries){
		double m1 = StatUtils.mean(xSeries);
		double m2 = StatUtils.mean(ySeries);
		double v1 = StatUtils.variance(xSeries, m1);
		double v2 = StatUtils.variance(ySeries, m2);
		double covariance = StatUtils.mean(xySeries) - (m1*m2);
			
		return (v1==0.0 || v2==0.0)?0.0: covariance/Math.pow(v1*v2,0.5);
	}
}
